package com.lalameow.researchstation.util.mixin

import com.lalameow.researchstation.Cbor
import com.lalameow.researchstation.Nbt
import com.lalameow.researchstation.ResearchStation
import com.lalameow.researchstation.api.researching.Researching
import com.lalameow.researchstation.api.technology.Technologies
import com.lalameow.researchstation.api.technology.Technology
import com.lalameow.researchstation.identifier
import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.builtins.SetSerializer
import kotlinx.serialization.decodeFromByteArray
import kotlinx.serialization.encodeToByteArray
import net.minecraft.entity.data.DataTracker
import net.minecraft.entity.data.TrackedDataHandler
import net.minecraft.entity.data.TrackedDataHandlerRegistry
import net.minecraft.entity.player.PlayerEntity
import net.minecraft.network.PacketByteBuf
import net.minecraft.util.Identifier
import net.silkmc.silk.core.serialization.serializers.ResourceLocationSerializer
import net.silkmc.silk.persistence.customCompoundKey
import net.silkmc.silk.persistence.persistentCompound

val PlayerEntity.researched: Set<Technology>
    get() = (this as Researching).researched

internal object ResearchingHelper {
    private val serializer = SetSerializer(ResourceLocationSerializer)
    private val key = customCompoundKey(
        ResearchStation.identifier("researching")!!,
        { ResearchStation.Nbt.encodeToNbtElement(serializer, it) },
        { ResearchStation.Nbt.decodeFromNbtElement(serializer, it) }
    )

    @OptIn(ExperimentalSerializationApi::class)
    private val RESEARCHING_DATA by lazy {
        object : TrackedDataHandler<MutableSet<Identifier>> {
            override fun write(buf: PacketByteBuf, value: MutableSet<Identifier>) {
                buf.writeByteArray(ResearchStation.Cbor.encodeToByteArray(value))
            }

            override fun read(buf: PacketByteBuf) =
                ResearchStation.Cbor.decodeFromByteArray<MutableSet<Identifier>>(buf.readByteArray())

            override fun copy(value: MutableSet<Identifier>) = value.toMutableSet()
        }
    }

    private val TECHNOLOGIES = DataTracker.registerData(PlayerEntity::class.java, RESEARCHING_DATA)!!

    init {
        TrackedDataHandlerRegistry.register(RESEARCHING_DATA)
    }

    fun PlayerEntity.initDataTracker() {
        dataTracker.startTracking(TECHNOLOGIES, mutableSetOf())
    }

    var PlayerEntity.technologies: MutableSet<Identifier>
        get() = dataTracker[TECHNOLOGIES]
        set(value) {
            dataTracker[TECHNOLOGIES] = value
        }

    val PlayerEntity.researched: Set<Technology>
        get() = technologies.mapNotNullTo(mutableSetOf()) { Technologies[it] }

    fun PlayerEntity.write() {
        persistentCompound[key] = technologies
    }

    fun PlayerEntity.read() {
        (this as MutableResearching).technologies = persistentCompound[key]?.toMutableSet() ?: mutableSetOf()
    }
}

internal interface MutableResearching : Researching {
    override var technologies: MutableSet<Identifier>
}
